# 导入必要的模块
import os
import uuid
from typing import List, Literal
from dotenv import load_dotenv
from langchain_core.messages import SystemMessage, AIMessage, HumanMessage, ToolMessage
from langchain_openai import ChatOpenAI
from langgraph.constants import END
from pydantic import BaseModel
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, MessageGraph

# 加载 .env 文件中的环境变量
load_dotenv()

# -------------------------------
# 定义系统模板字符串（系统提示词）
# -------------------------------
template = """你的任务是向用户收集信息，以了解他们想要创建什么类型的提示模板（prompt template）。

你需要从用户那里获取以下信息：

- 提示的目标（objective）是什么  
- 模板中会传入哪些变量（variables）  
- 对输出的限制（constraints），即输出中**不应出现或不应做的事情**  
- 对输出的要求（requirements），即输出中**必须遵守或必须包含的内容**

如果你无法确定这些信息，请向用户进一步确认！  
不要随意猜测。  
当你收集到全部必要信息后，调用相应的工具（tool）来继续执行。
"""


# -------------------------------
# 定义一个函数：将系统提示 + 用户消息 组合为消息列表
# -------------------------------
def get_messages_info(messages):
    return [SystemMessage(content=template)] + messages


# -------------------------------
# 定义数据模型（Pydantic）：用于存储提示模板的信息
# -------------------------------
class PromptInstructions(BaseModel):
    """提示模板的核心信息"""
    objective: str       # 目标
    variables: List[str] # 模板变量
    constraints: List[str] # 限制条件
    requirements: List[str] # 必须遵守的要求


# -------------------------------
# 初始化 LLM（连接 DeepSeek 模型）
# -------------------------------
llm = ChatOpenAI(
    api_key=os.getenv("MODELSCOPE_API_KEY"),
    openai_api_base=os.getenv("MODELSCOPE_API_BASE"),
    model="deepseek-ai/DeepSeek-R1-0528",
    streaming=True,
    temperature=0
)

# -------------------------------
# 绑定工具（PromptInstructions）
# -------------------------------
"""
在 LangChain / LangGraph 的工具机制中， “工具（tool）”不仅可以是函数，也可以是一个带参数结构的模型类（如 Pydantic 模型）。
这类“工具”其实是给大模型一个结构化输出的模板，
告诉它：“你需要输出的数据，必须符合这个类定义的格式。”
"""
llm_with_tool = llm.bind_tools([PromptInstructions])

# 定义第一阶段消息处理链
chain = get_messages_info | llm_with_tool


# -------------------------------
# 第二阶段系统提示模板：生成最终的提示模板
# -------------------------------
prompt_system = """请根据以下需求，编写一个高质量的提示模板：

{reqs}"""


# -------------------------------
# 定义函数：从消息中提取工具调用参数，生成提示模板消息
# -------------------------------
def get_prompt_messages(messages: list):
    tool_call = None
    other_msgs = []
    for m in messages:
        if isinstance(m, AIMessage) and m.tool_calls:
            tool_call = m.tool_calls[0]["args"]
        elif isinstance(m, ToolMessage):
            continue
        elif tool_call is not None:
            other_msgs.append(m)
    return [SystemMessage(content=prompt_system.format(reqs=tool_call))] + other_msgs


# 定义第二阶段消息处理链
prompt_gen_chain = get_prompt_messages | llm


# -------------------------------
# 状态判断函数（决定下一步流程）
# -------------------------------
def get_state(messages) -> Literal["add_tool_message", "info", "__end__"]:
    if isinstance(messages[-1], AIMessage) and messages[-1].tool_calls:
        return "add_tool_message"
    elif not isinstance(messages[-1], HumanMessage):
        return END
    return "info"


# -------------------------------
# 初始化 MemorySaver（内存型检查点保存器）
# -------------------------------
memory = MemorySaver()

# 初始化工作流（消息图）
workflow = MessageGraph()

# 添加节点
workflow.add_node("info", chain)
workflow.add_node("prompt_gen_chain", prompt_gen_chain)

# -------------------------------
# 定义一个节点：添加工具消息
# -------------------------------
@workflow.add_node
def add_tool_message(state: list):
    return ToolMessage(
        content="Prompt 已生成！",
        tool_call_id=state[-1].tool_calls[0]["id"]
    )


# -------------------------------
# 定义工作流的连接逻辑
# -------------------------------
workflow.add_edge(START, "info")
workflow.add_conditional_edges("info", get_state)
workflow.add_edge("add_tool_message", "prompt_gen_chain")
workflow.add_edge("prompt_gen_chain", END)

# -------------------------------
# 编译图，启用内存检查点
# -------------------------------
graph = workflow.compile(checkpointer=memory)


# -------------------------------
# 启动交互循环
# -------------------------------
config = {"configurable": {"thread_id": str(uuid.uuid4())}}

while True:
    user = input("用户输入 (输入 q/Q 退出): ")
    if user in {"q", "Q"}:
        print("AI: 再见～")
        break
    output = None
    # 将用户输入送入工作流
    for output in graph.stream(
        [HumanMessage(content=user)],
        config=config,
        stream_mode="updates"
    ):
        last_message = next(iter(output.values()))
        last_message.pretty_print()

    # 当生成最终提示时，打印“完成”
    if output and "prompt" in output:
        print("生成完成！")




'''
测试问题：

目标:收集客户满意度反馈。
变量:['客户名称'、'互动日期'、"提供的服务’、"评级(1-5 级)'、"评论"]
约束:['输出不应包含客户的任何个人身份信息(PII)。"]
要求:[输出必须包含结构化格式，其中包含上述每个变量的字段。"]


'''